import paddle
import paddle.nn.functional as F
import numpy as np
from load_data import load_data
from lenet import LeNetModel


# 定义模型训练设置
def train(model):
    model.train()
    # 使用SGD优化器，learning_rate设置为0.01
    opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
    # 训练5轮
    EPOCH_NUM = 5

    for epoch_id in range(EPOCH_NUM):
        for batch_id, data in enumerate(load_data()()):
            # 准备数据
            images,labels = data
            images = paddle.to_tensor(images)
            labels = paddle.to_tensor(labels)

            # 前向计算的过程
            predicts = model(images)

            # 计算损失, 取一个批次样本损失的平均值
            loss = F.cross_entropy(predicts, labels)
            avg_loss = paddle.mean(loss)

            # 每训练200批次的数据，输出当前损失的情况
            if batch_id % 200 == 0:
                print("epoch: {}, batch: {}, loss: {}".format(epoch_id, batch_id, avg_loss.numpy()))

            # 后向传播, 更新参数的过程
            avg_loss.backward()
            # 最小化损失，更新参数
            opt.step()
            # 清除梯度
            opt.clear_grad()

    paddle.save(model.state_dict(), '2-cnn/mnist-cnn.pdparams')

